import json
from flask import Flask, render_template, request, jsonify
from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings
import chromadb
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import os

app = Flask(__name__)

# 配置 NVIDIA API
os.environ["NVIDIA_API_KEY"] = "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"
# 初始化模型和 DB
llm = ChatNVIDIA(model="meta/llama3-70b-instruct", temperature=0.7, max_tokens=256)  # 降低 max_tokens
embedder = NVIDIAEmbeddings(model="NV-Embed-QA")
client = chromadb.PersistentClient("./chroma_db3")
collection = client.get_collection("text_collection3")

def is_daily_language(text: str) -> bool:
    # 假設用一些關鍵詞來驗證是否為日常用語
    daily_phrases = [
        # 问候与寒暄
        "你好", "Hi",
        "嗨", "Hey",
        "嘿", "Hello",
        "早上好", "Good morning",
        "下午好", "Good afternoon",
        "晚上好", "Good evening",
        "你今天过得怎么样？", "How are you today?",
        "我很好，谢谢！你呢？", "I'm fine, thank you! And you?",
        "还行，挺好的", "Not bad, pretty good",
        "不错，挺开心的", "Not bad, feeling good",
        "再见", "Goodbye",
        "拜拜", "Bye",
        "下次再聊", "Talk to you next time",
        "有空再联系", "Catch you later",
        "晚安", "Good night",

        # 请求与帮助
        "能帮我个忙吗？", "Can you help me?",
        "请问你能帮我看看这个吗？", "Could you help me with this?",
        "帮我一下呗！", "Give me a hand, please!",
        "当然可以！", "Sure!",
        "没问题，我来试试", "No problem, I'll give it a try",
        "好的，我尽力", "Okay, I'll do my best",

        # 表达感谢与道歉
        "谢谢你", "Thank you",
        "太感谢了", "Thank you so much",
        "真的帮了大忙", "You've been a great help",
        "对不起", "I'm sorry",
        "抱歉", "Sorry",
        "我错了，对不起", "I was wrong, I'm sorry",
        "真的很抱歉，我没想到", "I'm really sorry, I didn't mean to",

        # 表达情绪
        "我好开心", "I'm so happy",
        "太高兴了", "So glad",
        "哈哈，太好玩了", "Haha, that's so fun",
        "真的太棒了", "That's amazing",
        "唉，我好难过", "Oh, I'm so sad",
        "真的好沮丧", "I'm really down",
        "我心情不太好", "I'm not in a good mood",
        "我很生气", "I'm angry",
        "太气人了", "That's so frustrating",
        "这太让人受不了了", "This is too much",

        # 表达意愿与计划
        "我想……", "I want to...",
        "我想试试", "I'd like to try",
        "我打算……", "I plan to...",
        "我计划……", "I'm planning to...",
        "我想去……", "I want to go to...",
        "我想要……", "I want...",
        "你想去吗？", "Do you want to go?",
        "你打算怎么安排？", "How are you planning?",
        "你想要什么？", "What do you want?",

        # 日常交流
        "你最近忙不忙？", "Are you busy lately?",
        "你今天做了什么？", "What did you do today?",
        "你有没有空？", "Are you free?",
        "我今天挺忙的", "I'm pretty busy today",
        "我刚做完……", "I just finished...",
        "我现在有点累", "I'm a bit tired now",

        # 表达同意与不同意
        "好的", "Okay",
        "行", "Alright",
        "没问题", "No problem",
        "我同意", "I agree",
        "我觉得也是", "I think so too",
        "没错，你说得对", "You're right",
        "我不太同意", "I don't quite agree",
        "我觉得不太一样", "I think it's different",
        "我不太确定", "I'm not sure",

        # 表达疑问
        "是吗？", "Really?",
        "真的吗？", "Is that true?",
        "怎么回事？", "What's going on?",
        "为什么？", "Why?",
        "你能再说清楚一点吗？", "Can you explain that again?",
        "这是怎么回事？", "What's the deal with this?",
        "我不太明白，能解释一下吗？", "I don't quite understand, can you explain?",

        # 轻松幽默的表达
        "你这是在逗我呢吧？", "Are you kidding me?",
        "哈哈，你太搞笑了", "Haha, you're so funny",
        "这太有意思了", "This is so interesting",
        "我真是个笨蛋", "I'm such an idiot",
        "我又犯傻了", "I did it again",
        "我真是服了我自己", "I give up on myself",

        # 表达不确定
        "我猜……", "I guess...",
        "我觉得可能是……", "I think it might be...",
        "我不太确定，大概是……", "I'm not sure, probably...",
        "你能再说清楚一点吗？", "Can you be more specific?",
        "这是这样吗？", "Is that so?",
        "我没听错吧？", "Did I hear that right?",
        "这个你是意思吗？", "Is that what you mean?",

        # 表达惊讶
        "哇，太厉害了", "Wow, that's amazing",
        "天哪，我没想到", "Oh my god, I didn't expect that",
        "真的假的？", "Are you serious?",

        # 表达安慰
        "别难过，一切都会好起来的", "Don't be sad, everything will be okay",
        "没关系，别太在意", "It's okay, don't worry about it",
        "别担心，我陪着你", "Don't worry, I'm here with you",

        # 确认
        "好的，我知道了", "Okay, I got it",
        "明白了", "Understood",
        "嗯，我懂了", "Yeah, I get it",

        # 拒绝
        "我不太方便", "I'm not really comfortable with that",
        "我可能不行", "I don't think I can",
        "我不太能接受", "I'm not quite ready for that"
    ]
    return any(phrase in text for phrase in daily_phrases)

def truncate_text(text, max_tokens=7936):
    """截断文本以确保不超过 max_tokens"""
    tokens = text.split()
    if len(tokens) > max_tokens:
        tokens = tokens[:max_tokens]
    return " ".join(tokens)

def rag_answer(question: str) -> str:
    # 检索逻辑
    results = collection.query(
        query_texts=[question],
        n_results=3,  # 减少检索结果数量
        include=["documents", "metadatas"]
    )

    # 检查 results 是否包含有效数据
    if not results or not results["documents"]:
        return "未找到相关上下文信息。"

    context = []
    for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
        # 确保 meta 不为 None，并提供默认值
        source = meta.get('source', '未知') if meta else '未知'
        context.append(f"{doc} (来源: {source})")

    # 截断上下文以确保总长度不超过 7936 tokens
    context_text = "\n".join(context)
    max_context_tokens = 7936 - len(question.split())  # 保留问题长度
    truncated_context = truncate_text(context_text, max_tokens=max_context_tokens)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "基于上下文的内容请使用中文回答，需注明来源:"),
        ("user", "上下文:\n{context}\n\n问题: {input}")
    ])

    chain = prompt | llm | StrOutputParser()
    answer = chain.invoke({"input": question, "context": "\n".join(context)})

    # 根據回答的情況構建 JSON 輸出
    if is_daily_language(question):
        result_json = json.dumps({"msg": "hello", "content": "None"})
    elif results["documents"]:  # 如果有匹配的結果
        result_json = answer
    else:
        result_json = json.dumps({"msg": "未找到相關信息", "content": "None"})

    # # 将字典转换为 JSON 字符串
    # result_str = json.dumps(result_json)

    # 打印 JSON 字符串
    print(result_json)

    return result_json

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/ask', methods=['POST'])
def ask_question():
    data = request.json
    try:
        answer = rag_answer(data['question'])
        return jsonify({"answer": answer})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(host='127.0.0.1', port=5200, debug=True)